import os
from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms


def cum_sum(a: list):
    tmp = 0
    result = []
    for i in a:
        tmp = tmp + i
        result.append(tmp)
    return result


class WSIdatasets(Dataset):

    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels = os.listdir(data_dir)
        nums = [len(os.listdir(os.path.join(data_dir, i))) for i in self.labels]
        nums.insert(0, 0)
        self.start_index = cum_sum(nums)
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(norm_mean, norm_std)
        ])

    def __getitem__(self, index):
        length = len(self.start_index)
        for i in range(length):
            if index < self.start_index[i + 1]:
                y = self.labels[i]
                cur_path = os.path.join(self.data_dir, y)
                cur_ind = index - self.start_index[i]
                cur_data_path = os.path.join(cur_path, os.listdir(cur_path)[cur_ind])
                x = Image.open(cur_data_path).convert('RGB')
                x = self.transform(x)
                y = i
                return x, y
        return None

    def __len__(self):
        return self.start_index[-1]

def build_wsi_dataset(data_dir):
    return WSIdatasets(data_dir)